#pragma once

#include <vector>
#include <stdint.h>
#include <string>

class FashionMnistLoad
{
public:
    FashionMnistLoad(const std::string& resource_dir);

    int64_t load();

    const std::vector<std::vector<uint8_t>>& get_train_images();
    const std::vector<uint8_t>& get_train_labels();

    const std::vector<uint8_t>& get_test_labels();
    const std::vector<std::vector<uint8_t>>& get_test_images();

    const std::vector<std::vector<uint8_t>>& get_one_hot_train_labels();
    const std::vector<std::vector<uint8_t>>& get_one_hot_test_labels();

    const std::vector<std::vector<double>>& get_normalized_train_images();
    const std::vector<std::vector<double>>& get_normalized_test_images();
private:
    std::vector<std::vector<uint8_t>> load_images(const std::string& path, int max_items = -1);
    std::vector<uint8_t> load_labels(const std::string& path, int max_items = -1);
    std::vector<std::vector<double>> normalize(const std::vector<std::vector<uint8_t>> &input);
private:
    std::string resource_dir_{};
    std::vector<std::vector<uint8_t>> train_images_{};
    std::vector<std::vector<double>> normalized_train_images_{};
    std::vector<uint8_t> train_labels_{};
    std::vector<std::vector<uint8_t>> one_hot_train_labels_{};
    std::vector<std::vector<uint8_t>> test_images_{};
    std::vector<std::vector<double>> normalized_test_images_{};
    std::vector<uint8_t> test_labels_{};
    std::vector<std::vector<uint8_t>> one_hot_test_labels_{};
};
